#include "../src/planner/query_planner.h"
#include "../src/parser/sql_parser.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <inttypes.h>

// Test helper functions
static void print_test_header(const char* test_name);
static void print_test_result(const char* test_name, bool passed);

static void print_test_header(const char* test_name) {
    printf("\n=== %s ===\n", test_name);
}

static void print_test_result(const char* test_name, bool passed) {
    printf("%s: %s\n", test_name, passed ? "PASSED" : "FAILED");
}

// Test basic planner creation and destruction
static bool test_planner_basic(void) {
    print_test_header("Testing Basic Planner Operations");
    
    query_planner_t* planner = planner_create();
    if (!planner) {
        printf("Failed to create planner\n");
        return false;
    }
    
    printf("Created planner successfully\n");
    printf("Optimization enabled: %s\n", planner->enable_optimization ? "true" : "false");
    printf("Optimization level: %u\n", planner->optimization_level);
    
    planner_destroy(planner);
    printf("Destroyed planner successfully\n");
    
    return true;
}

// Test plan node creation
static bool test_plan_node_creation(void) {
    print_test_header("Testing Plan Node Creation");
    
    bool all_passed = true;
    
    // Test scan plan creation
    plan_node_t* scan_plan = planner_create_scan_plan("users", NULL);
    if (!scan_plan) {
        printf("Failed to create scan plan\n");
        all_passed = false;
    } else {
        printf("Created scan plan for table 'users'\n");
        printf("  Node type: %s\n", plan_node_type_to_string(scan_plan->type));
        printf("  Table name: %s\n", scan_plan->data.scan.table_name);
        printf("  Estimated rows: %" PRIu64 "\n", scan_plan->data.scan.stats.row_count);
    }
    
    // Test index scan plan creation
    plan_node_t* index_scan_plan = planner_create_index_scan_plan("users", "idx_name", NULL);
    if (!index_scan_plan) {
        printf("Failed to create index scan plan\n");
        all_passed = false;
    } else {
        printf("Created index scan plan\n");
        printf("  Node type: %s\n", plan_node_type_to_string(index_scan_plan->type));
        printf("  Table name: %s\n", index_scan_plan->data.index_scan.table_name);
        printf("  Index name: %s\n", index_scan_plan->data.index_scan.index_name);
    }
    
    // Test join plan creation
    if (scan_plan && index_scan_plan) {
        plan_node_t* join_plan = planner_create_join_plan(JOIN_INNER, NULL, scan_plan, index_scan_plan);
        if (!join_plan) {
            printf("Failed to create join plan\n");
            all_passed = false;
        } else {
            printf("Created join plan\n");
            printf("  Node type: %s\n", plan_node_type_to_string(join_plan->type));
            printf("  Join type: %s\n", join_type_to_string(join_plan->data.join.join_type));
            printf("  Child count: %u\n", join_plan->child_count);
            
            plan_destroy(join_plan); // This will also destroy children
            scan_plan = NULL;
            index_scan_plan = NULL;
        }
    }
    
    // Clean up remaining nodes
    if (scan_plan) plan_destroy(scan_plan);
    if (index_scan_plan) plan_destroy(index_scan_plan);
    
    return all_passed;
}

// Test cost estimation
static bool test_cost_estimation(void) {
    print_test_header("Testing Cost Estimation");
    
    query_planner_t* planner = planner_create();
    if (!planner) {
        printf("Failed to create planner\n");
        return false;
    }
    
    bool all_passed = true;
    
    // Create a simple scan plan
    plan_node_t* scan_plan = planner_create_scan_plan("test_table", NULL);
    if (!scan_plan) {
        printf("Failed to create scan plan\n");
        planner_destroy(planner);
        return false;
    }
    
    // Estimate cost
    plan_cost_t cost = estimate_plan_cost(scan_plan, planner->optimizer);
    
    printf("Cost estimation for scan plan:\n");
    printf("  CPU cost: %.2f\n", cost.cpu_cost);
    printf("  I/O cost: %.2f\n", cost.io_cost);
    printf("  Memory cost: %.2f\n", cost.memory_cost);
    printf("  Total cost: %.2f\n", cost.total_cost);
    printf("  Cardinality: %" PRIu64 "\n", cost.cardinality);
    
    if (cost.total_cost <= 0) {
        printf("Invalid cost estimation\n");
        all_passed = false;
    }
    
    plan_destroy(scan_plan);
    planner_destroy(planner);
    
    return all_passed;
}

// Test memo group management
static bool test_memo_management(void) {
    print_test_header("Testing Memo Management");
    
    optimizer_context_t* ctx = optimizer_create_context();
    if (!ctx) {
        printf("Failed to create optimizer context\n");
        return false;
    }
    
    bool all_passed = true;
    
    // Create memo groups
    memo_group_t* group1 = memo_create_group(ctx);
    memo_group_t* group2 = memo_create_group(ctx);
    
    if (!group1 || !group2) {
        printf("Failed to create memo groups\n");
        all_passed = false;
    } else {
        printf("Created memo groups:\n");
        printf("  Group 1 ID: %u\n", group1->group_id);
        printf("  Group 2 ID: %u\n", group2->group_id);
        
        // Test group lookup
        memo_group_t* found = memo_find_group(ctx, group1->group_id);
        if (found != group1) {
            printf("Failed to find memo group\n");
            all_passed = false;
        } else {
            printf("Successfully found memo group\n");
        }
    }
    
    printf("Total memo groups: %u\n", ctx->memo_group_count);
    
    optimizer_destroy_context(ctx);
    
    return all_passed;
}

// Test basic optimization rules
static bool test_optimization_rules(void) {
    print_test_header("Testing Optimization Rules");
    
    bool all_passed = true;
    
    // Test rule conditions
    plan_node_t* join_plan = planner_create_join_plan(JOIN_INNER, NULL, NULL, NULL);
    if (!join_plan) {
        printf("Failed to create join plan for rule testing\n");
        return false;
    }
    
    // Test join commutativity rule
    bool can_apply = rule_join_commutativity_condition(join_plan);
    printf("Join commutativity rule applicable: %s\n", can_apply ? "true" : "false");
    
    if (can_apply) {
        optimizer_context_t* ctx = optimizer_create_context();
        if (ctx) {
            plan_node_t* transformed = rule_join_commutativity_transform(join_plan, ctx);
            if (transformed) {
                printf("Successfully applied join commutativity rule\n");
                plan_destroy(transformed);
            } else {
                printf("Failed to apply join commutativity rule\n");
                all_passed = false;
            }
            optimizer_destroy_context(ctx);
        }
    }
    
    plan_destroy(join_plan);
    
    return all_passed;
}

// Test plan tree printing
static bool test_plan_printing(void) {
    print_test_header("Testing Plan Tree Printing");
    
    // Create a complex plan tree
    plan_node_t* scan1 = planner_create_scan_plan("table1", NULL);
    plan_node_t* scan2 = planner_create_scan_plan("table2", NULL);
    
    if (!scan1 || !scan2) {
        printf("Failed to create scan plans\n");
        return false;
    }
    
    plan_node_t* join = planner_create_join_plan(JOIN_INNER, NULL, scan1, scan2);
    if (!join) {
        printf("Failed to create join plan\n");
        plan_destroy(scan1);
        plan_destroy(scan2);
        return false;
    }
    
    // Create projection on top of join
    expression_t* dummy_expr = NULL;  // We don't have actual expressions yet
    char* col_name = "result";
    plan_node_t* project = planner_create_project_plan(&dummy_expr, &col_name, 1, join);
    
    if (!project) {
        printf("Failed to create project plan\n");
        plan_destroy(join);
        return false;
    }
    
    printf("Plan tree structure:\n");
    plan_print_tree(project, 0);
    
    plan_destroy(project);
    
    return true;
}

// Test integration with SQL parser
static bool test_sql_integration(void) {
    print_test_header("Testing SQL Parser Integration");
    
    query_planner_t* planner = planner_create();
    if (!planner) {
        printf("Failed to create planner\n");
        return false;
    }
    
    // Create a simple SQL statement
    const char* sql = "SELECT name FROM users WHERE age > 18";
    printf("Parsing SQL: %s\n", sql);
    
    parser_t* parser = parser_create(sql);
    if (!parser) {
        printf("Failed to create SQL parser\n");
        planner_destroy(planner);
        return false;
    }
    
    statement_t* stmt = parser_parse_statement(parser);
    bool success = false;
    
    if (parser_has_error(parser)) {
        printf("SQL parsing error: %s\n", parser_get_error(parser));
    } else if (stmt) {
        printf("Successfully parsed SQL statement\n");
        printf("Statement type: %s\n", stmt_type_to_string(stmt->type));
        
        // Try to create a plan from the statement
        plan_node_t* plan = planner_create_plan(planner, stmt);
        if (plan) {
            printf("Successfully created execution plan\n");
            printf("Plan tree:\n");
            plan_print_tree(plan, 0);
            
            // Print planner statistics
            planner_print_stats(planner);
            
            success = true;
        } else {
            printf("Failed to create execution plan\n");
        }
        
        statement_destroy(stmt);
    } else {
        printf("Failed to parse SQL statement\n");
    }
    
    parser_destroy(parser);
    planner_destroy(planner);
    
    return success;
}

int main() {
    printf("Query Planner Test Suite\n");
    printf("========================\n");
    
    int passed = 0;
    int total = 0;
    
    // Run all tests
    struct {
        const char* name;
        bool (*test_func)(void);
    } tests[] = {
        {"Basic Planner Operations", test_planner_basic},
        {"Plan Node Creation", test_plan_node_creation},
        {"Cost Estimation", test_cost_estimation},
        {"Memo Management", test_memo_management},
        {"Optimization Rules", test_optimization_rules},
        {"Plan Tree Printing", test_plan_printing},
        {"SQL Parser Integration", test_sql_integration},
        {NULL, NULL}
    };
    
    for (int i = 0; tests[i].name != NULL; i++) {
        total++;
        bool result = tests[i].test_func();
        if (result) {
            passed++;
        }
        print_test_result(tests[i].name, result);
    }
    
    printf("\n=== Test Summary ===\n");
    printf("Passed: %d/%d\n", passed, total);
    printf("Success rate: %.1f%%\n", (double)passed / total * 100.0);
    
    return (passed == total) ? 0 : 1;
}